#include "raytracer.hpp"
#include "image.hpp"

typedef std::list<GeometryNode*> GeometryList;
GeometryList m_primitives;
std::list<Light*> r_lights;
Colour r_ambient = Colour(0,0,0);
SceneNode* r_root;
Vector3D lt_dir; 	// direction to left top pixel
Vector3D dpi_down;	 
Vector3D dpi_right;
Point3D min_space;
Point3D max_space;

int ANTIALIASING;
int BAD_PIXEL;
int DEPTHOFFIELD;
double SHARP_DEPTH;
int UNIFORMDIVISION;
int GRIDNUM;

int dx, dy, dz;
double dx_reci, dy_reci, dz_reci, shiftx, shifty, shiftz;
GRID ***mySpace = NULL;
void a4_render(// What to render
               SceneNode* root,
               // Where to output the image
               const std::string& filename,
               // Image size
               int width, int height,
               // Viewing parameters
               const Point3D& eye, const Vector3D& view,
               const Vector3D& up, double fov,
               // Lighting parameters
               const Colour& ambient,
               const std::list<Light*>& lights,
				int anti_aliasing,
				int bad_pixel,
				int depthoffield,
				double sharp_depth,
				int subdivision,
				int gridnum
				)
{
	ANTIALIASING = anti_aliasing;
	BAD_PIXEL = bad_pixel;
	DEPTHOFFIELD = depthoffield;
	SHARP_DEPTH = sharp_depth;
	UNIFORMDIVISION = subdivision;
	GRIDNUM = gridnum;
	if(UNIFORMDIVISION){
        root->flat_hier(Matrix4x4());
        find_space(eye);
        build_space();
	}
	Image img(width, height, 3);
	// global variables setup
	r_lights = lights;
	r_ambient = ambient;
	r_root = root;
	// ray direction initialization
	Vector3D c_up = up;
	Point3D center = eye + view;
	double d = sqrt(view.dot(view));
	double factor = tan((fov*M_PI)/360);
	double h = factor * d * 2 ;
	double w = (double)width/height * h;
	Vector3D c_right = view.cross(up);
  	c_up.normalize();
  	c_right.normalize();
	center = center - (double)w/2*c_right + (double)h/2*c_up;
	lt_dir = center - eye;
	dpi_down = - (double)(h/height)*c_up;
	dpi_right = (double)(w/width)*c_right;
	
	int counter = 0;
	for (int y = 0; y < height; y++) {
		Vector3D tempV = lt_dir + (double)y*dpi_down;
		for (int x = 0; x < width; x++) {
			Vector3D V = tempV + (double)x*dpi_right;
			V.normalize();
			Ray m_ray(eye, V);
            Colour colour(0,0,0);
			if(DEPTHOFFIELD){
                Point3D sharp_point(0,0,SHARP_DEPTH);
                Vector3D N = Vector3D(0,0,1);
                double dot_product = N.dot(V);
                double u = -N.dot(eye - sharp_point)/dot_product;
                Point3D temp_point = eye + u*V;
                Colour blur_colour(0,0,0);
                for(int i = 0; i < 20; i++){
                    Point3D blureye = Point3D(eye[0] + 10*(double)rand()/RAND_MAX, eye[1] + 10*(double)rand()/RAND_MAX, eye[2]);
                    Vector3D blurdir = temp_point - blureye;
                    blurdir.normalize();
                    Ray blur_ray(blureye, blurdir);
                    blur_colour = blur_colour + raytrace(blur_ray,0);
                }
                colour = blur_colour/20;
            }
            else colour = raytrace(m_ray, 0);
            
			img(x,y,0)= colour.R();
			img(x,y,1)= colour.G();
			img(x,y,2)= colour.B();
        }
		int complete = (int)((double)(y+1)/height*100.0);
		if(complete == counter){ 
			std::cout<< complete << "% COMPLETED" << std::endl;
			counter += 10;	
		}  	
	}
	
if(ANTIALIASING){
    // adaptive anti-aliasing
	bool bad_pixel[width][height];
	for(int y = 0; y < height; y++){
		for(int x = 0; x < width; x++){
			if(x == 1 || x==width-1 || y==1 || y==height-1){
				bad_pixel[x][y] = false;
				continue;
			}
			Colour px1 = Colour(img(x-1,y-1,0),img(x-1,y-1,1),img(x-1,y-1,2));
			Colour px2 = Colour(img(x,y-1,0),img(x,y-1,1),img(x,y-1,2));
			Colour px3 = Colour(img(x+1,y-1,0),img(x+1,y-1,1),img(x+1,y-1,2));
			Colour px4 = Colour(img(x-1,y,0),img(x-1,y,1),img(x-1,y,2));
			Colour px5 = Colour(img(x,y,0),img(x,y,1),img(x,y,2));
			Colour px6 = Colour(img(x+1,y,0),img(x+1,y,1),img(x+1,y,2));
			Colour px7 = Colour(img(x-1,y+1,0),img(x-1,y+1,1),img(x-1,y+1,2));
			Colour px8 = Colour(img(x,y+1,0),img(x,y+1,1),img(x,y+1,2));
			Colour px9 = Colour(img(x+1,y+1,0),img(x+1,y+1,1),img(x+1,y+1,2));
			double diff = 0.1;
			if( px5-px1 > diff || px5-px2 > diff || px5-px3 > diff ||
				px5-px4 > diff || px5-px6 > diff ||
				px5-px7 > diff || px5-px8 > diff || px5-px9 > diff){
				bad_pixel[x][y] = true;
			}
			else{
				bad_pixel[x][y] = false;
			}
		}
	}
	// hightlight bad_pixel
	for(int y=0; y<height; y++){
		for(int x=0; x<width; x++){
			if(bad_pixel[x][y]){
if(BAD_PIXEL){
                    img(x,y,0) = 1.0; img(x,y,1) = 0.0; img(x,y,2) = 0.0;
}else{// NOT SHOW BAD_PIXEL
                // cast 9 rays and get average for the pixel
                    Colour sum(0.0,0.0,0.0);
                    for(int i=0; i<3; i++){
                        for(int j=0; j<3; j++){
                            Vector3D V = lt_dir + (x+0.5*(i-1))*dpi_right + (y+0.5*(j-1))*dpi_down;
                            V.normalize();
                            Ray m_ray(eye, V);
                            Colour temp_colour(0,0,0);
                            if(DEPTHOFFIELD){
                                Point3D sharp_point(0,0,SHARP_DEPTH);
                                Vector3D N = Vector3D(0,0,1);
                                double dot_product = N.dot(V);
                                double u = -N.dot(eye - sharp_point)/dot_product;
                                Point3D temp_point = eye + u*V;
                                Colour blur_colour(0,0,0);
                                for(int i = 0; i < 20; i++){
                                    Point3D blureye = Point3D(eye[0] + 10*(double)rand()/RAND_MAX, eye[1] + 10*(double)rand()/RAND_MAX, eye[2]);
                                    Vector3D blurdir = temp_point - blureye;
                                    blurdir.normalize();
                                    Ray blur_ray(blureye, blurdir);
                                    blur_colour = blur_colour + raytrace(blur_ray,0);
                                }
                                temp_colour = blur_colour/20;
                            }
                            else temp_colour = raytrace(m_ray, 0);
                            sum = sum + temp_colour;
                        }
                    }
                    sum = sum/9.0;
                    img(x,y,0)= sum.R();
                    img(x,y,1)= sum.G();
                    img(x,y,2)= sum.B();
}// NOT SHOW BAD_PIXEL
            }
        }
    }
}// ANTIALIASING
  	img.savePng(filename);
}

Colour raytrace(Ray ray, int depth)
{
	if(depth > 15) return Colour(0,0,0);
	Intersection m_intersection;
	Colour colour(0,0,0);
	Colour refr_colour(0,0,0);
	Colour refl_colour(0,0,0);
	Colour phong(0,0,0);
	bool hit = false;
	if(UNIFORMDIVISION){
		Point3D pos = ray.get_position();
		Vector3D dir = ray.get_direction();
		int i = int((pos[0]+shiftx)*dx_reci);
		int j = int((pos[1]+shifty)*dy_reci);
		int k = int((pos[2]+shiftz)*dz_reci);
		if(i < 0) i = 0; if(i > GRIDNUM-1) i = GRIDNUM-1; 
		if(j < 0) j = 0; if(j > GRIDNUM-1) j = GRIDNUM-1;
		if(k < 0) k = 0; if(k > GRIDNUM-1) k = GRIDNUM-1;
		GRID m_grid = mySpace[i][j][k];
		int stepX = (dir[0]<0)? -1:1;
		int stepY = (dir[1]<0)? -1:1;
		int stepZ = (dir[2]<0)? -1:1;
		Point3D px = m_grid.min_point, py = px, pz = px;
		if(stepX == 1){
			px = Point3D(px[0]+dx, px[1], px[2]);
		}
		if(stepY == 1){
			py = Point3D(py[0], py[1]+dy, py[2]);
		}
		if(stepZ == 1){
			pz = Point3D(pz[0], pz[1], pz[2]+dz);
		}
		Vector3D normalX(stepX,0,0); double dpX = normalX.dot(dir);
		Vector3D normalY(0,stepY,0); double dpY = normalY.dot(dir);
		Vector3D normalZ(0,0,stepZ); double dpZ = normalZ.dot(dir);
		double tMaxX = dpX==0? INT_MAX: -normalX.dot(pos-px)/dpX;
		double tMaxY = dpY==0? INT_MAX: -normalY.dot(pos-py)/dpY;
		double tMaxZ = dpZ==0? INT_MAX: -normalZ.dot(pos-pz)/dpZ;
		double tDeltaX = dir[0]==0 ? 0:fabs(dx/dir[0]);
		double tDeltaY = dir[1]==0 ? 0:fabs(dy/dir[1]);
		double tDeltaZ = dir[2]==0 ? 0:fabs(dz/dir[2]);
		double MAX_DIST = INT_MAX;
		bool id_table[256] = {false};
		for(;;){
			std::list<GeometryNode*>::const_iterator node;
			for(node = m_grid.geolist.begin(); node!=m_grid.geolist.end(); node++){
				int id = (*node)->get_id();
				if(id_table[id]) { continue; }
				else{id_table[id] = true;}
				bool temp_hit = (*node)->ray_intersect(ray, &m_intersection);
				if( temp_hit && m_intersection.m_dist < MAX_DIST){
					MAX_DIST = m_intersection.m_dist;
				}
				if( temp_hit ) hit = true;
			}
			if(tMaxX <= tMaxY && tMaxX <= tMaxZ){
				i += stepX;
				if( i < 0 || i >= GRIDNUM) break;
				tMaxX += tDeltaX;
			}
			else if(tMaxY <= tMaxX && tMaxY <= tMaxZ){
				j += stepY;
				if( j < 0 || j >= GRIDNUM) break;
				tMaxY += tDeltaY;
			}
			else if(tMaxZ <= tMaxY && tMaxZ <= tMaxX){
				k += stepZ;
				if( k < 0 || k >= GRIDNUM) break;
				tMaxZ += tDeltaZ;
			}
			m_grid = mySpace[i][j][k];
		}
	}
	else{
		//if(UNIFORMDIVISION) hit = cast_ray_to_primitives(ray, &m_intersection);
		 hit = r_root->ray_intersect(ray, &m_intersection); 
	}
	if(hit){
		Material *material = m_intersection.m_material;
		double mir_coef = material->get_reflection();
		double refract = material->get_refraction();
		
		// phong lighting
		if(refract == 0){
			phong = do_phong_reflection(m_intersection, ray);
		}
		
		// reflection
		if(mir_coef > 0.0){
			Vector3D N = m_intersection.m_normal;
			Vector3D V = ray.get_direction();
			N.normalize();
			V.normalize();
			Vector3D R(V - 2.0 * V.dot(N) * N);
			R.normalize();
			Ray reflect_ray(m_intersection.m_intersect, R);
			refl_colour = raytrace(reflect_ray, depth+1);
			Colour ks = m_intersection.m_material->get_specular();
			refl_colour = refl_colour*ks;
		}
		// refraction
		if(refract > 0.0){
			Vector3D R(0,0,0);
			Vector3D N = m_intersection.m_normal;
			Vector3D V = ray.get_direction();
			V.normalize();
			N.normalize();
			Vector3D NL = N.dot(V)<0? N:-N;
			bool into = N.dot(NL) > 0;
			double nc=1, nt = refract;
			double nnt = into? nc/nt: nt/nc;
			double ddn = V.dot(NL);
			double cos2t = 1 - pow(nnt,2)*(1-pow(ddn,2));
			if(cos2t < 0){
				R = V - 2.0*N.dot(V)*N;
				R.normalize();
			}
			else{
				R = nnt*V-(into?1:-1)*(ddn*nnt+sqrt(cos2t))*N;
				R.normalize();
			}
			
			Ray refract_ray(m_intersection.m_intersect, R);
			Colour absorbance = m_intersection.m_material->get_specular()*0.20*(-1.0);
			Colour transparency = Colour(exp(absorbance.R()),exp(absorbance.G()),exp(absorbance.B()));
			refr_colour = transparency*raytrace(refract_ray, depth+1);
		}
	
		colour = phong + refr_colour + mir_coef*refl_colour;
	}
	return Colour(	std::min(1.0, colour.R()),
					std::min(1.0, colour.G()),
					std::min(1.0, colour.B()));
}

/* Illumination = ambient + SUM( diffuse*(light_ray.normal)*lightintensity + 
 *							specular(reflect.view_ray)^shininess*lightintensity )
 * reflect = 2(light_ray.normal)*normal - light_ray
 * all vectors are normalized
 */
Colour do_phong_reflection(const Intersection& intersection, Ray ray)
{
	Texture *texture = intersection.m_texture;
	Material *material = intersection.m_material;
	Colour colour(0,0,0);
	Colour m_diffuse = material->get_diffuse();
	if(texture){
		m_diffuse = texture->extract_colour(intersection.u, intersection.v);
	}
	Vector3D hit_normal = intersection.m_normal;
	hit_normal.normalize();
	Colour m_specular = material->get_specular();
	double m_shininess = material->get_shininess();
	Point3D hit_pos = intersection.m_intersect;
	Vector3D view_ray = - ray.get_direction();
	view_ray.normalize();
	colour = colour + r_ambient*m_diffuse;
	for (std::list<Light*>::const_iterator I = r_lights.begin(); I != r_lights.end(); ++I) {
		Point3D light_pos = (*I)->position;
		Vector3D light_ray = light_pos - hit_pos;
		light_ray.normalize();
		int num_hit = 0;
		int num_light = (*I)->num_light;
		if( num_light < 1) {
			std::cout << "light number should be at least 1" << std::endl;
			exit(-1);
		}
        if( num_light > 1 ){
            int light_diff = (*I)->dist;
            double x = light_pos[0]-num_light/2*light_diff;
            double y = light_pos[1]-num_light/2*light_diff;
            double z = light_pos[2];
            for(int i=0; i<num_light; i++){
                for(int j=0; j<num_light; j++){
                    Intersection shadow_intersection;
                    Point3D distributed_pos(x+i*light_diff,y+i*light_diff,z);
                    Vector3D distributed_ray(distributed_pos - hit_pos);
                    distributed_ray.normalize();
                    Ray light(hit_pos, distributed_ray);
                    bool shadow = false;
                    if(UNIFORMDIVISION)
                        shadow = cast_ray_to_primitives(light, &shadow_intersection);
                    else 
                        shadow = r_root->ray_intersect(light, &shadow_intersection);
                    if(shadow) num_hit++;
                }
            }
        	if(num_hit != pow(num_light,2) ){
				Vector3D reflect = 2.0f * (light_ray.dot(hit_normal)) * hit_normal - light_ray;
				reflect.normalize();
            	Colour phong(0,0,0);
				if(light_ray.dot(hit_normal) > 0)
					phong = phong + m_diffuse*(light_ray.dot(hit_normal))*(*I)->colour;
#ifdef BLIN
				Vector3D H = (view_ray+light_ray);
				H.normalize();
				if(H.dot(hit_normal) > 0)
					phong = phong + m_specular*pow(H.dot(hit_normal), m_shininess)*(*I)->colour;
#else
				if(reflect.dot(view_ray) > 0)
					phong = phong + m_specular*pow(reflect.dot(view_ray), m_shininess)*(*I)->colour;
#endif 
            	double factor = 1.0 - (double)num_hit/pow(num_light,2);
            	colour = colour + Colour(phong.R()*factor, phong.G()*factor, phong.B()*factor);
			}
        }
        else{
            Ray light(hit_pos, light_ray);
            bool shadow = false;
            Intersection shadow_intersection;
            if(UNIFORMDIVISION)
                shadow = cast_ray_to_primitives(light, &shadow_intersection);
            else 
                shadow = r_root->ray_intersect(light, &shadow_intersection);
            if(!shadow){
                Vector3D reflect = 2.0f * (light_ray.dot(hit_normal)) * hit_normal - light_ray;
                reflect.normalize();
                Colour phong(0,0,0);
                if(light_ray.dot(hit_normal) > 0)
                    phong = phong + m_diffuse*(light_ray.dot(hit_normal))*(*I)->colour;
#ifdef BLIN
                Vector3D H = (view_ray+light_ray);
                H.normalize();
                if(H.dot(hit_normal) > 0)
                    phong = phong + m_specular*pow(H.dot(hit_normal), m_shininess)*(*I)->colour;
#else
                if(reflect.dot(view_ray) > 0)
                    phong = phong + m_specular*pow(reflect.dot(view_ray), m_shininess)*(*I)->colour;
#endif 
                colour = colour + phong;
            }
        }
	}
	return colour;
}

bool cast_ray_to_primitives(Ray ray, Intersection *intersection)
{
	double MAX_DIST = 1000000.0f;
	bool hit = false;
	for (GeometryList::const_iterator node = m_primitives.begin(); node != m_primitives.end(); node++) {
		bool temp_hit = (*node)->ray_intersect(ray, intersection);
		if( temp_hit && intersection->m_dist < MAX_DIST){
			MAX_DIST = intersection->m_dist;
		}
		if( temp_hit ) hit = true;
	}
	return hit;
}

void find_space(const Point3D& eye)
{
	min_space = Point3D(INT_MAX,INT_MAX,INT_MAX);
	max_space = Point3D(INT_MIN,INT_MIN,INT_MIN);
	for (GeometryList::const_iterator node = m_primitives.begin(); node != m_primitives.end(); node++) {
		std::pair<Point3D,Point3D> result = (*node)->get_voxels();
		Point3D temp_min = result.first;
		Point3D temp_max = result.second;
		min_space[0] = std::min(min_space[0],temp_min[0]);
		min_space[1] = std::min(min_space[1],temp_min[1]);
		min_space[2] = std::min(min_space[2],temp_min[2]);
		max_space[0] = std::max(max_space[0],temp_max[0]);
		max_space[1] = std::max(max_space[1],temp_max[1]);
		max_space[2] = std::max(max_space[2],temp_max[2]);
	}
	min_space[0] = (std::min(min_space[0],eye[0]));
	min_space[1] = (std::min(min_space[1],eye[1]));
	min_space[2] = (std::min(min_space[2],eye[2]));
	max_space[0] = (std::max(max_space[0],eye[0]));
	max_space[1] = (std::max(max_space[1],eye[1]));
	max_space[2] = (std::max(max_space[2],eye[2]));
	// extend space
	min_space[0] = floor(min_space[0]/10)*10-10;
	min_space[1] = floor(min_space[1]/10)*10-10;
	min_space[2] = floor(min_space[2]/10)*10-10;
	max_space[0] = ceil(max_space[0]/10)*10+10;
	max_space[1] = ceil(max_space[1]/10)*10+10;
	max_space[2] = ceil(max_space[2]/10)*10+10;
}


void build_space()
{
	dx = ((double)fabs(min_space[0] - max_space[0])/GRIDNUM);
	dy = ((double)fabs(min_space[1] - max_space[1])/GRIDNUM);
	dz = ((double)fabs(min_space[2] - max_space[2])/GRIDNUM);
	// refine grid edges
	dx = (ceil(dx/10)+1)*10;
	dy = (ceil(dy/10)+1)*10;
	dz = (ceil(dz/10)+1)*10;
	dx_reci = 1.0f/dx; dy_reci = 1.0f/dy; dz_reci = 1.0f/dz;
	shiftx = -min_space[0];
	shifty = -min_space[1];
	shiftz = -min_space[2];
	mySpace = new GRID**[GRIDNUM];
	for(int i=0; i<GRIDNUM; i++){
		mySpace[i] = new GRID*[GRIDNUM];
		for(int j=0; j<GRIDNUM; j++){
			mySpace[i][j] = new GRID[GRIDNUM];
			for(int k=0; k<GRIDNUM; k++){
				mySpace[i][j][k] = 
					GRID(Point3D(min_space[0]+i*dx,min_space[1]+j*dy,min_space[2]+k*dz),
					Point3D(min_space[0]+(i+1)*dx,min_space[1]+(j+1)*dy,min_space[2]+(k+1)*dz));
			}
		}
	}
	for (GeometryList::const_iterator node = m_primitives.begin(); node != m_primitives.end(); node++) {
		std::pair<Point3D,Point3D> temp_pair = (*node)->get_voxels();
		Point3D min_point = temp_pair.first;
		Point3D max_point = temp_pair.second;
		int xi = (int)((min_point[0]+shiftx)*dx_reci);
		int xf = (int)((max_point[0]+shiftx)*dx_reci)+1;
		if(xi < 0) xi = 0; if(xi >= GRIDNUM) xi = GRIDNUM;
		if(xf < 0) xf = 0; if(xf >= GRIDNUM) xf = GRIDNUM; 
		int yi = (int)((min_point[1]+shifty)*dy_reci);
		int yf = (int)((max_point[1]+shifty)*dy_reci)+1;
		if(yi < 0) xi = 0; if(yi >= GRIDNUM) yi = GRIDNUM;
		if(yf < 0) xf = 0; if(yf >= GRIDNUM) yf = GRIDNUM;
		int zi = (int)((min_point[2]+shiftz)*dz_reci);
		int zf = (int)((max_point[2]+shiftz)*dz_reci)+1;
		if(zi < 0) zi = 0; if(zi >= GRIDNUM) zi = GRIDNUM;
		if(zf < 0) zf = 0; if(zf >= GRIDNUM) zf = GRIDNUM;
		for(int i=xi; i<=xf && i<GRIDNUM; i++){
			for(int j=yi; j<=yf && j<GRIDNUM; j++){
				for(int k=zi; k<=zf && k<GRIDNUM; k++){
					mySpace[i][j][k].geolist.push_back((GeometryNode*)(*node));
				}
			}
		}
	}
}